#include "skinPainterContext.h"

#include <maya/MGlobal.h>
#include <maya/MUIDrawManager.h>
#include <maya/MFnMesh.h>
#include <maya/MFrameContext.h>
#include <maya/MToolsInfo.h>
#include <maya/MHWGeometry.h>
#include <maya/MItMeshPolygon.h>
#include <maya/MItMeshVertex.h>
#include <maya/MPointArray.h>
#include <maya/MTransformationMatrix.h>

#include <QtGui/QKeyEvent>
#include <QtCore/QCoreApplication>

#include "customSkinDeformer.h"


MString skinPainterContext::toolTitle = "XS Skin Brush";
MString skinPainterContext::toolHelpText = "Draw weights what xsSkinDeformer";


skinPainterContext::skinPainterContext()
{
	setTitleString(skinPainterContext::toolTitle);
	setHelpString(skinPainterContext::toolHelpText);


	// init attr
	brush.brushSize = 1.0f;
	brush.brushColor = MColor(1.0f, 0.0f, 0.0f);
	brush.brushWidth = 2.0f;

	brushMode = BrushMode::kNormal;
	//currentJointIndex = 0;
	numInfluenceJoints = 0;
	paintOperation = PaintOperation::kAdd;
}

skinPainterContext::~skinPainterContext()
{
}

void skinPainterContext::toolOnSetup(MEvent& event)
{
	MPxContext::toolOnSetup(event);

	QCoreApplication* app = qApp;
	app->installEventFilter(this);
}

void skinPainterContext::toolOffCleanup()
{
	QCoreApplication* app = qApp;
	app->removeEventFilter(this);
	
	MPxContext::toolOffCleanup();
}

MStatus skinPainterContext::doPress(MEvent& event, MHWRender::MUIDrawManager& drawMgr,
                                    const MHWRender::MFrameContext& context)
{
	short cursorX, cursorY;
	event.getPosition(cursorX, cursorY);
	lastCursorPos = MPoint(cursorX, cursorY);
	currentCursorPos = MPoint(cursorX, cursorY);
	bUpdateBrushPos = true;

	DrawEffects(event, drawMgr, context);

	if (brushMode == BrushMode::kNormal)
		PaintWeights();
	
	return MStatus::kSuccess;
}

MStatus skinPainterContext::doRelease(MEvent& event, MHWRender::MUIDrawManager& drawMgr,
	const MHWRender::MFrameContext& context)
{
	bUpdateBrushPos = true;

	DrawEffects(event, drawMgr, context);
	
	return MStatus::kSuccess;
}

MStatus skinPainterContext::doPtrMoved(MEvent& event, MHWRender::MUIDrawManager& drawMgr,
                                       const MHWRender::MFrameContext& context)
{
	view3D = M3dView::active3dView();
	
	DrawEffects(event, drawMgr, context);
	return MStatus::kSuccess;
}

MStatus skinPainterContext::doDrag(MEvent& event, MHWRender::MUIDrawManager& drawMgr,
                                   const MHWRender::MFrameContext& context)
{

	short cursorX, cursorY;
	event.getPosition(cursorX, cursorY);
	lastCursorPos = currentCursorPos;
	currentCursorPos = MPoint(cursorX, cursorY);
	
	if (brushMode == BrushMode::kBrushSize)
	{
		bUpdateBrushPos = false;
		
		float scaleSize = currentCursorPos.distanceTo(lastCursorPos) * 0.05f;
		if (currentCursorPos[0] > lastCursorPos[0])
		{
			SetBrushSize(BrushSize() + scaleSize);
		}
		else
		{
			SetBrushSize(std::max(BrushSize() - scaleSize, 0.01f));
		}
	}
	else
	{
		PaintWeights();
	}

	
	DrawEffects(event, drawMgr, context);
	
	return MStatus::kSuccess;
}

bool skinPainterContext::eventFilter(QObject* watched, QEvent* event)
{
	if (QKeyEvent* keyEvent = dynamic_cast<QKeyEvent*>(event))
	{
		if (keyEvent->type() == QEvent::KeyPress)
		{
			doKeyPress(keyEvent);
		}
		else if(keyEvent->type() == QEvent::KeyRelease)
		{
			doKeyRelease(keyEvent);
		}
	}
	return QObject::eventFilter(watched, event);
}

void skinPainterContext::SetBrushColor(double r, double g, double b)
{
	brush.brushColor = MColor(r,g,b);
	MToolsInfo::setDirtyFlag(*this);
}

void skinPainterContext::SetBrushOpacity(double opacity)
{
	brush.opacity = opacity;
	MToolsInfo::setDirtyFlag(*this);
}

void skinPainterContext::SetBrushValue(double val)
{
	brush.value = val;
	MToolsInfo::setDirtyFlag(*this);
}

void skinPainterContext::SetDrawMesh(const MDagPath& dagpathMesh)
{
	drawMesh = dagpathMesh;
	MToolsInfo::setDirtyFlag(*this);
}

void skinPainterContext::doKeyPress(QKeyEvent* keyEvent)
{
	if (keyEvent->key() == Qt::Key_B)
	{
		brushMode = BrushMode::kBrushSize;
	}
}

void skinPainterContext::doKeyRelease(QKeyEvent* keyEvent)
{
	brushMode = BrushMode::kNormal;
}

void skinPainterContext::UpdateBrushPos(MEvent & event, const MHWRender::MFrameContext& context)
{
	// 获取当前鼠标世界位置
	short cursorX, cursorY;
	MPoint cursorPosNear, cursorPosFar;
	event.getPosition(cursorX, cursorY);
	context.viewportToWorld(cursorX, cursorY, cursorPosNear, cursorPosFar);

	MVector cursorVecWorld = cursorPosFar - cursorPosNear;
	cursorVecWorld.normalize();

	// 射线检测
	MFloatPoint hitPoint;
	MFnMesh fmnMesh(drawMesh);
	brush.isHit = fmnMesh.closestIntersection(cursorPosNear, cursorVecWorld, nullptr, nullptr, false, MSpace::kWorld, 9999, false,
								nullptr, hitPoint, nullptr, nullptr, nullptr, nullptr, nullptr);

	brush.brushPos = hitPoint;
	fmnMesh.getClosestNormal(brush.brushPos, brush.brushNormal, MSpace::kWorld);
	brush.brushNormal.normalize();

	// update brush mat
	MVector worldUP(0.0, 1.0, 0.0);
	MVector axisY = brush.brushNormal;
	MVector axisZ;
	MVector axisX;
	
	if ((axisY * worldUP) == 1.0)
	{
		axisZ = MVector(0.0, 0.0, 1.0);
		axisX = MVector(1.0, 0.0, 0.0);
	}
	else
	{
		axisZ = axisY ^ worldUP;
		axisZ.normalize();
		axisX = axisY ^ axisZ;
		axisX.normalize();
	}

	double srcMat[4][4] = {
		{axisX.x, axisX.y, axisX.z, 0.0},
		{axisY.x, axisY.y, axisY.z, 0.0},
		{axisZ.x, axisZ.y, axisZ.z, 0.0},
		{brush.brushPos.x, brush.brushPos.y, brush.brushPos.z, 1.0}};
	
	brush.worldMat = MMatrix(srcMat);
}

void skinPainterContext::DrawBrush(MEvent & event, MHWRender::MUIDrawManager& drawMgr, const MHWRender::MFrameContext& context)
{
	if (bUpdateBrushPos)
		UpdateBrushPos(event, context);
	
	if (!brush.isHit)
		return;

	MPointArray linePoints;
	linePoints.append(MPoint(-5, 0, -5));
	linePoints.append(MPoint(5, 0, -5));
	linePoints.append(MPoint(5, 0, 5));
	linePoints.append(MPoint(-5, 0, 5));
	
	for (auto& point:linePoints)
	{
		point *= brush.worldMat;
	}
	
	drawMgr.beginDrawable();
	{
		
		drawMgr.setColor(brush.brushColor);
	
		// draw point
		drawMgr.setPointSize(8.f);
		drawMgr.point(brush.brushPos);
	
		// draw circle
		drawMgr.setLineWidth(brush.brushWidth);
		drawMgr.circle(brush.brushPos, brush.brushNormal, brush.brushSize, false);
	}
	drawMgr.endDrawable();
}

// void skinPainterContext::DrawMeshWireframe(MHWRender::MUIDrawManager& drawMgr)
// {
// 	if (!drawMesh.isValid())
// 		return;
// 	
// 	drawMgr.beginDrawable();
// 	{
// 		drawMgr.setDepthPriority(3);
// 		drawMgr.setColor(MColor(1.0,1.0,1.0));
// 		drawMgr.setLineWidth(1.5f);
//
// 		MPointArray points;
// 		for (auto wireIndices:wireframeIndices)
// 			drawMgr.mesh(MUIDrawManager::kLineStrip, meshVertices, &meshVtxNormals, nullptr, &wireIndices);
// 	}
// 	drawMgr.endDrawable();
// }

// void skinPainterContext::DrawMesh(MHWRender::MUIDrawManager& drawMgr)
// {
// 	UpdateVtxColors();
// 	drawMgr.beginDrawable();
// 	{
// 		drawMgr.setDepthPriority(5);
// 		drawMgr.mesh(MUIDrawManager::kTriangles, meshVertices, &meshVtxNormals, &vtxColors, &triangleIndoces);
// 		drawMgr.mesh(MUIDrawManager::kTriStrip, meshVertices, &meshVtxNormals, &vtxColors, &triangleStripIndoces);
// 	}
// 	drawMgr.endDrawable();
// }

void skinPainterContext::SetBrushSize(float size)
{
	brush.brushSize = size;
	MToolsInfo::setDirtyFlag(*this);
}

float skinPainterContext::BrushSize() const
{
	return brush.brushSize;
}

void skinPainterContext::SetVtxIndices(const MIntArray& inTriangle, const MIntArray& inTriangleStrip,
	const std::vector<MIntArray> inWireframe)
{
	triangleIndoces.clear();
	for (auto idx:inTriangle)
		triangleIndoces.append(idx);

	triangleStripIndoces.clear();
	for (auto idx:inTriangleStrip)
		triangleStripIndoces.append(idx);

	wireframeIndices.clear();
	for (auto indicesBlock:inWireframe)
	{
		MUintArray wireIndices;
		for (auto idx:indicesBlock)
			wireIndices.append(idx);

		wireframeIndices.emplace_back(wireIndices);
	}

	//UpdateMeshPoints();
}

void skinPainterContext::SetPaintJointIndex(unsigned int jointIdx)
{
	currentPaintJointIndex = jointIdx;
	//UpdateVtxColors();
	MToolsInfo::setDirtyFlag(*this);
}

// void skinPainterContext::UpdateVtxColors()
// {
// 	vtxColors.clear();
// 	MFnDependencyNode mfnDPMesh(drawMesh.node());
// 	MObject xsSkinObj = mfnDPMesh.findPlug("inMesh", false).source().node();
//
// 	MFnDependencyNode mfnDPSkin(xsSkinObj);
// 	MPlug plugWeightList = mfnDPSkin.findPlug("weightList", false);
// 	
// 	unsigned int pointsNum = plugWeightList.numElements();
// 	for(unsigned int i = 0; i < pointsNum; i++)
// 	{
// 		MPlug weights = plugWeightList.elementByLogicalIndex(i).child(0);
// 		unsigned int numInfluenceJnts = weights.numElements();
// 		unsigned int currentJntIdx = currentPaintJointIndex < numInfluenceJnts ? currentPaintJointIndex : (numInfluenceJnts - 1);
//
// 		MPlug weight_i = weights.elementByLogicalIndex(currentJntIdx);
// 		double weight = weight_i.asDouble();
// 		MColor vtxColor(MColor::kRGB, weight, 0, 0);
// 		vtxColors.append(vtxColor);
// 	}
// 	
// 	MToolsInfo::setDirtyFlag(*this);
// }

void skinPainterContext::UpdateInfluenceJnt(const MIntArray& influenceJnts)
{
	currentUnlockJointIndices.clear();
	currentUnlockJointIndices.copy(influenceJnts);

	MToolsInfo::setDirtyFlag(*this);
}

void skinPainterContext::SetInfluenceJointNum(unsigned int jntNum)
{
	numInfluenceJoints = jntNum;
	
	MToolsInfo::setDirtyFlag(*this);
}

void skinPainterContext::SetPaintOperation(const PaintOperation& op)
{
	paintOperation = op;
}

// void skinPainterContext::UpdateMeshPoints()
// {
// 	if (!drawMesh.isValid())
// 		return;
//
// 	meshVertices.clear();
// 	meshVtxNormals.clear();
//
// 	MFnMesh mfnMesh(drawMesh);
// 	
// 	mfnMesh.getPoints(meshVertices, MSpace::kWorld);
// 	mfnMesh.getNormals(meshVtxNormals, MSpace::kWorld);
//
// 	MToolsInfo::setDirtyFlag(*this);
// }

void skinPainterContext::GetPointsInBrushRange(MIntArray& outVtxIndexInRange)
{
	if (!drawMesh.isValid())
		return;
	
	MMatrix brushWorldMatInverse = brush.worldMat.inverse();
	float brushRadius = brush.brushSize;
	
	MItMeshPolygon mitMeshPoly(drawMesh);
	for (; !mitMeshPoly.isDone(); mitMeshPoly.next())
	{
		// exclusive back face
		MVector faceNormal;
		mitMeshPoly.getNormal(faceNormal, MSpace::kWorld);
		if ((faceNormal * brush.brushNormal) <= 0.0)
			continue;
		
		MPointArray points;
		mitMeshPoly.getPoints(points, MSpace::kWorld);

		MIntArray indices;
		mitMeshPoly.getVertices(indices);

		for (unsigned int i = 0; i < indices.length(); ++i)
		{
			MPoint vtxPos = points[i] * brushWorldMatInverse;
			if (vtxPos.distanceTo(MPoint(0.0,0.0,0.0)) < brushRadius)
			{
				// avoid contain repeatedly
				int vtxIndex = indices[i];
				bool isExisted = false;
				for (auto& idx:outVtxIndexInRange)
				{
					if (vtxIndex == idx)
					{
						isExisted = true;
						break;
					}
				}

				if (!isExisted)
					outVtxIndexInRange.append(vtxIndex);
			}
		}
	}
}

void skinPainterContext::DrawWeightPoints(const MIntArray& inVtxIndices, MHWRender::MUIDrawManager& drawMgr)
{
	if (!(drawMesh.isValid() && brush.isHit))
		return;
	
	MFnMesh mfnMesh(drawMesh);

	drawMgr.beginDrawable();
	{
		drawMgr.setPointSize(3.0);
		drawMgr.setColor(brush.brushColor);
		
		for (auto idx:inVtxIndices)
		{
			MPoint vtx;
			mfnMesh.getPoint(idx, vtx, MSpace::kWorld);
			drawMgr.point(vtx);
		}
	}
	drawMgr.endDrawable();

}

void skinPainterContext::DrawEffects(MEvent & event, MHWRender::MUIDrawManager& drawMgr, const MHWRender::MFrameContext& context)
{
	DrawBrush(event, drawMgr, context);

	MIntArray vtxIndices;
	GetPointsInBrushRange(vtxIndices);

	DrawWeightPoints(vtxIndices, drawMgr);
}

void skinPainterContext::GetNeighborPoints(int findPoint, MIntArray& outNeighborPoints)
{
	MItMeshVertex mitMeshVtx(drawMesh);
	for (; !mitMeshVtx.isDone(); mitMeshVtx.next())
	{
		if (mitMeshVtx.index() == findPoint)
		{
			mitMeshVtx.getConnectedVertices(outNeighborPoints);
			break;
		}
	}
}

double skinPainterContext::GetVtxWeight(int vtxIndex, int jointIndex)
{
	MFnDependencyNode mfndpMesh(drawMesh.node());
	MPlug plugInMesh = mfndpMesh.findPlug("inMesh", false);
	MFnDependencyNode mfndpDeformer(plugInMesh.source().node());

	MPlug plugWeightList = mfndpDeformer.findPlug("weightList", false);
	MPlug weightList_i = plugWeightList.elementByLogicalIndex(vtxIndex);
	MPlug plugWeights = weightList_i.child(0).elementByLogicalIndex(jointIndex);

	return plugWeights.asDouble();
}

void skinPainterContext::SetVtxWeight(int vtxIndex, int jointIndex, double weight)
{
	MIntArray otherJointIndices;
	for (auto jntIdx:currentUnlockJointIndices)
	{
		if (jntIdx != jointIndex)
		{
			otherJointIndices.append(jntIdx);
		}
	}
	unsigned int numOtherJnt = otherJointIndices.length();
	if (numOtherJnt == 0)
		return;

	weight = weight > 1.0 ? 1.0 : weight;
	
	MFnDependencyNode mfndpMesh(drawMesh.node());
	MPlug plugInMesh = mfndpMesh.findPlug("inMesh", false);
	MFnDependencyNode mfndpDeformer(plugInMesh.source().node());
	
	MPlug plugWeightList = mfndpDeformer.findPlug("weightList", false);
	MPlug weightList_i = plugWeightList.elementByLogicalIndex(vtxIndex);

	MPlug plugWeights = weightList_i.child(0);

	MPlug plugWeight_i = plugWeights.elementByLogicalIndex(jointIndex);
	double prevWeight = plugWeight_i.asDouble();
	double actualWeight = prevWeight;

	// ensure weights count == 1.0
	double delta = prevWeight - weight;
	if (delta < 0)
	{
		double deltaTotal = abs(delta);
		double averageWeight = deltaTotal / numOtherJnt;
		for (auto otherIdx:otherJointIndices)
		{
			plugWeight_i = plugWeights.elementByLogicalIndex(otherIdx);
			prevWeight = plugWeight_i.asDouble();
			double newWeight = 0.0;

			numOtherJnt -= 1;
			if (prevWeight < averageWeight)
			{
				deltaTotal -= prevWeight;
				actualWeight += prevWeight;

				if (numOtherJnt > 0)
					averageWeight = deltaTotal / numOtherJnt;
			}
			else
			{
				newWeight = prevWeight - averageWeight;
				actualWeight += averageWeight;
			}
			plugWeight_i.setDouble(newWeight);
		}
	}
	else
	{
		actualWeight -= delta;
		double averageWeight = delta / numOtherJnt;
		for (auto otherIdx:otherJointIndices)
		{
			plugWeight_i = plugWeights.elementByLogicalIndex(otherIdx);
			prevWeight = plugWeight_i.asDouble();
			plugWeight_i.setDouble(prevWeight + averageWeight);
		}
	}

	plugWeight_i = plugWeights.elementByLogicalIndex(jointIndex);
	plugWeight_i.setDouble(actualWeight);
	
	// normalize weights
	double weightTotal = 0.0;
	MDoubleArray normalizeWeights;
	double currentWeight = 0.0;
	for (unsigned int i = 0; i < numInfluenceJoints; ++i)
	{
		plugWeight_i = plugWeights.elementByLogicalIndex(i);
		currentWeight = plugWeight_i.asDouble();
		weightTotal += currentWeight;
		normalizeWeights.append(currentWeight);
	}
	
	for (unsigned int i = 0; i < numInfluenceJoints; ++i)
	{
		plugWeight_i = plugWeights.elementByLogicalIndex(i);
		plugWeight_i.setDouble(normalizeWeights[i] / weightTotal);
	}
}

void skinPainterContext::RefreshDeformer()
{
	MFnDependencyNode mfndpMesh(drawMesh.node());
	MPlug plugInMesh = mfndpMesh.findPlug("inMesh", false);
	MFnDependencyNode mfndpDeformer(plugInMesh.source().node());
	MPlug plugAffetcOutputGeomAttr= mfndpDeformer.findPlug("affetcOutputGeomAttr", false);
	plugAffetcOutputGeomAttr.setInt(0);
}

void skinPainterContext::PaintWeights()
{
	if (IsJntLocked(currentPaintJointIndex))
		return;
	
	switch (paintOperation)
	{
	case PaintOperation::kAdd:
		AddWeight();
		break;
	case PaintOperation::kReplace:
		ReplaceWeight();
		break;
	case PaintOperation::kSmooth:
		SmoothWeight();
		break;
	}
}

bool skinPainterContext::IsJntLocked(int jntIdx)
{
	bool result = true;
	for (auto& idx:currentUnlockJointIndices)
	{
		if (jntIdx == idx)
		{
			result = false;
			break;
		}
	}

	return result;
}

void skinPainterContext::AddWeight()
{
	MIntArray vtxIndices;
	GetPointsInBrushRange(vtxIndices);
	// vertex index
	for (auto vtxIdx:vtxIndices)
	{
		double currentW = GetVtxWeight(vtxIdx, currentPaintJointIndex);
		double actualW = Lerp(currentW, currentW + brush.value, brush.opacity);
		SetVtxWeight(vtxIdx, currentPaintJointIndex, actualW);
	}

	RefreshDeformer();
	view3D.refresh(true);
}

void skinPainterContext::ReplaceWeight()
{
	MIntArray vtxIndices;
	GetPointsInBrushRange(vtxIndices);
	// vertex index
	for (auto vtxIdx:vtxIndices)
	{
		double currentW = GetVtxWeight(vtxIdx, currentPaintJointIndex);
		double actualW = Lerp(currentW, brush.value, brush.opacity);
		SetVtxWeight(vtxIdx, currentPaintJointIndex, actualW);
	}

	RefreshDeformer();
	view3D.refresh(true);
}

void skinPainterContext::SmoothWeight()
{
	MIntArray vtxIndices;
	GetPointsInBrushRange(vtxIndices);

	// vertex index
	for (auto vtxIdx:vtxIndices)
	{
		// calculate weight
		// influence joint index
		for (auto jntIdx:currentUnlockJointIndices)
		{
			MIntArray neighborPoints;
			// Get neighbor vertex
			GetNeighborPoints(vtxIdx, neighborPoints);
			
			// Get all weight for nrighbor vtx about jnt
			double neighborVtxWeights = 0.0;
			for (auto neighborVtxIndex:neighborPoints)
			{
				neighborVtxWeights += GetVtxWeight(neighborVtxIndex, jntIdx);
			}
			// average weight
			double neighborAverageWeight = neighborVtxWeights / neighborPoints.length();
			double currentWeight = GetVtxWeight(vtxIdx, jntIdx);
			double actualW = Lerp(currentWeight, neighborAverageWeight, brush.opacity, 0.1);
			SetVtxWeight(vtxIdx, jntIdx, actualW);
		}
	}

	RefreshDeformer();
	view3D.refresh(true);
}

double skinPainterContext::Lerp(double start, double end, double ratio, double scale)
{
	return (end - start) * ratio * scale + start;
}

